// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <optional>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::test;

namespace {

class ArrayNormalizeTest : public FunctionBaseTest {
 protected:
  void testExpr(
      const VectorPtr& expected,
      const std::vector<VectorPtr>& input) {
    auto result = evaluate("array_normalize(C0, C1)", make_row_vector(input));
    assertEqualVectors(expected, result);
  }

  // Helper function to calculate p-norm based on its definition
  // and normalize the input array.
  template <typename T>
  std::vector<T> normalize(const std::vector<T>& values, T p) {
    T sum = 0;
    for (const auto& item : values) {
      sum += pow(abs(item), p);
    }

    T pNorm = pow(sum, 1.0 / p);

    std::vector<T> res;
    for (T value : values) {
      res.push_back(value / pNorm);
    }

    return res;
  }

  template <typename T>
  void testArrayWithElementsEqualZero() {
    auto input = make_array_vector<T>(
        {{static_cast<T>(0.0), static_cast<T>(-0.0), static_cast<T>(0.0)}});
    auto p = make_constant<T>(static_cast<T>(2.0), input->size());
    testExpr(input, {input, p});
  }

  template <typename T>
  void testArrayWithPLessThanZero() {
    auto input = make_array_vector<T>(
        {{static_cast<T>(1.0), static_cast<T>(2.0), static_cast<T>(3.0)}});
    auto p = make_constant<T>(static_cast<T>(-1.0), input->size());
    POLLUX_ASSERT_THROW(
        testExpr(input, {input, p}),
        "array_normalize only supports non-negative p");
  }

  template <typename T>
  void testArrayWithPEqualZero() {
    auto vector = make_array_vector<T>(
        {{static_cast<T>(1.0), static_cast<T>(-2.0), static_cast<T>(3.0)}});
    auto p = make_constant<T>(static_cast<T>(0.0), vector->size());
    testExpr(vector, {vector, p});
  }

  template <typename T>
  void testArrayWithPLessThanOne() {
    T pValue = 0.5;
    std::vector<T> values = {1.0, -4.0, 9.0};
    std::vector<T> expected = normalize(values, pValue);

    auto input = make_array_vector<T>({values});
    auto p = make_constant<T>(pValue, input->size());
    testExpr(make_array_vector<T>({expected}), {input, p});
  }

  template <typename T>
  void testArrayWithPEqualOne() {
    T pValue = 1.0;
    std::vector<T> values = {1, -2, 3};
    std::vector<T> expected = normalize(values, pValue);

    auto input = make_array_vector<T>({values});
    auto p = make_constant<T>(pValue, input->size());
    testExpr(make_array_vector<T>({expected}), {input, p});
  }

  template <typename T>
  void testArrayWithTypeLimits() {
    auto input = make_array_vector<T>({
        {},
        {0},
        {1},
        {std::numeric_limits<T>::infinity()},
        {std::numeric_limits<T>::quiet_NaN()},
        {std::numeric_limits<T>::lowest(), -1.0},
        {std::numeric_limits<T>::max(), 1.0},
    });
    auto p = make_constant<T>(2.0, input->size());
    auto expected = make_array_vector<T>(
        {{},
         {0},
         {1},
         {std::numeric_limits<T>::quiet_NaN()},
         {std::numeric_limits<T>::quiet_NaN()},
         // For the cases below, pNorm is calculated as Infinity, resulting in
         // zero values. The result matches the Presto Java version.
         {-0.0, -0.0},
         {0.0, 0.0}});
    testExpr(expected, {input, p});
  }

  template <typename T>
  void testArrayWithNullValues() {
    auto input = make_nullable_array_vector<T>(
        {{{1, 2, std::nullopt}}, {{std::nullopt}}, std::nullopt});
    auto p = make_constant<T>(2.0, input->size());
    auto expected =
        make_nullable_array_vector<T>({std::nullopt, std::nullopt, std::nullopt});
    testExpr(expected, {input, p});
  }

  template <typename T>
  void testArrayWithDifferentValues() {
    T pValue = 4.0;

    std::vector<std::vector<T>> inputs = {
        {1.0},
        {-1.0, -1.0},
        {1.0, 4.0, 9.0},
        {0.1, 0.01, 1.0, 0.0},
        {2.34, 4.56, 8.12},
        {-2.34, 4.56, -8.12}};

    std::vector<std::vector<T>> expectedResults;
    for (const auto& input : inputs) {
      expectedResults.push_back(normalize(input, pValue));
    }

    auto input = make_array_vector<T>(inputs);
    auto p = make_constant<T>(pValue, input->size());
    // The test results were also calculated manually and checked against
    // Presto Java version.
    testExpr(make_array_vector<T>(expectedResults), {input, p});
  }
};

TEST_F(ArrayNormalizeTest, arrayWithElementsZero) {
  testArrayWithElementsEqualZero<float>();
  testArrayWithElementsEqualZero<double>();
}

TEST_F(ArrayNormalizeTest, pLessThanZero) {
  testArrayWithPLessThanZero<float>();
  testArrayWithPLessThanZero<double>();
}

TEST_F(ArrayNormalizeTest, pEqualsZero) {
  testArrayWithPEqualZero<float>();
  testArrayWithPEqualZero<double>();
}

TEST_F(ArrayNormalizeTest, pLessThanOne) {
  testArrayWithPLessThanOne<float>();
  testArrayWithPLessThanOne<double>();
}

TEST_F(ArrayNormalizeTest, pEqualsOne) {
  testArrayWithPEqualOne<float>();
  testArrayWithPEqualOne<double>();
}

TEST_F(ArrayNormalizeTest, limits) {
  testArrayWithTypeLimits<float>();
  testArrayWithTypeLimits<double>();
}

TEST_F(ArrayNormalizeTest, nullValues) {
  testArrayWithNullValues<float>();
  testArrayWithNullValues<double>();
}

TEST_F(ArrayNormalizeTest, differentValues) {
  testArrayWithDifferentValues<float>();
  testArrayWithDifferentValues<double>();
}

TEST_F(ArrayNormalizeTest, throwForIntTypes) {
  std::string errorMessage =
      "array_normalize only supports double and float types";
  POLLUX_ASSERT_THROW(testArrayWithElementsEqualZero<int>(), errorMessage);
  POLLUX_ASSERT_THROW(testArrayWithPEqualZero<int>(), errorMessage);
  POLLUX_ASSERT_THROW(testArrayWithPEqualOne<int>(), errorMessage);
}

} // namespace
